from iteration_helpers import update
import networkx as nx
import pandas as pd
from broker_helpers import add_optimal_broker, optimal_broker_update, add_social_broker, social_broker_update, add_structural_broker, structural_broker_update, add_power_broker, power_broker_update

def run_iterations_with_broker(G, total_iterations=5000, record_interval=25,
                               broker_type='N/A', model_type='N/A'):
    assortativity_values = []
    iterations = []
    broker_name = 200
    broker_added = False

    for i in range(total_iterations):

        # Every `record_interval` iterations, record the assortativity coefficient
        if (i + 1) % record_interval == 0:
            assortativity = nx.attribute_assortativity_coefficient(G, 'identity')
            assortativity_values.append(assortativity)
            iterations.append(i + 1)
            if not broker_added and assortativity > 0.6:
                if broker_type == 'optimal':
                  add_optimal_broker(G, broker_name)
                  print("Optimal Broker added at iteration " + str(i + 1))
                  broker_added = True
                if broker_type == 'social':
                  add_social_broker(G, broker_name)
                  print("Social Broker added at iteration " + str(i + 1))
                  broker_added = True
                if broker_type == 'structural':
                  add_structural_broker(G, broker_name)
                  print("Structural Broker added at iteration " + str(i + 1))
                  broker_added = True
                if broker_type == 'power':
                  add_power_broker(G, broker_name)
                  print("Power Broker added at iteration " + str(i + 1))
                  broker_added = True
                   
        
        if broker_added:
            if broker_type.startswith('optimal'):
              optimal_broker_update(G, broker_name)
            if broker_type.startswith('social'):
              social_broker_update(G, broker_name)
            if broker_type.startswith('structural'):
              structural_broker_update(G, broker_name)
            if broker_type.startswith('power'):
              power_broker_update(G, broker_name)
        
        else:
            update(G)


    assort_df = pd.DataFrame({
        'Iteration': iterations,
        'Assortativity': assortativity_values
    })
    filename = f"{broker_type}_{model_type}.csv"
    assort_df.to_csv(filename, index=False)